"""
Copyright (c) 2024 Beijing Jiaotong University
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Yunzhou Tang
Created: 2024/03/06
Supported by: National Key Research and Development Program of China
"""

import phot

if __name__ == '__main__':
    """ 双偏振光收发模块 + 光纤G652 """

    phot.config(plot=True, backend="torch")

    """ 系统参数 """
    num_symbols = 2 ** 16  # 符号数目
    bits_per_symbol = 4  # 2 for QPSK; 4 for 16QAM; 6 for 64QAM  设置调制格式
    total_baud = 20e9  # 信号波特率，符号率
    up_sampling_factor = 4  # 上采样倍数
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    Reference_frequency = 193.1e12

    """ 首先产生发射端X/Y双偏振信号 """
    signal_bits = phot.gen_bits(num_symbols * bits_per_symbol)

    """ 生成双偏振符号序列 """
    signals = phot.modulation(signal_bits, bits_per_symbol)

    prev_symbols = signals

    """ 上采样 """
    signals = phot.up_sample(signals, up_sampling_factor)

    """ RRC """
    RRC_ROLL_OFF = 0.01
    signals = phot.pulse_shaper(
        signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    """ DAC """
    DAC_Resolution_Bits = 10  # DAC量化bit位数
    DAC_OUTPUT_VPP = 1
    signals = phot.dac(signals, DAC_Resolution_Bits, DAC_OUTPUT_VPP)

    """ 光源 """
    Linewidth = 100e3
    output_power_dB = 0
    laser_center_frequency = Reference_frequency
    Optical_tx_laser = phot.laser_cw(num_symbols * up_sampling_factor, laser_center_frequency, Reference_frequency,
                                     sampling_rate, output_power_dB, Linewidth)

    """ 调制器 """
    Extinction_ratio = 60
    VPI = 5
    VDC = -2.5
    signals = phot.iq_modulator(
        signals, Optical_tx_laser, VDC, Extinction_ratio, VPI)

    # """ 发射端电滤波器 """
    # filter_type = 'Gaussian'
    # BW_electrical = total_baud * 0.75
    # order = 4
    # signals = phot.Electrical_filter_POL(
    #     signals, sampling_rate, filter_type, BW_electrical, order)

    """ OSNR """
    osnr = 30  # 设置系统OSNR，也就是光信号功率与噪声功率的比值，此处单位为dB
    signals = phot.gaussian_noise(signals, osnr, sampling_rate)

    """ Optical Fiber Channel """
    num_spans = 2  # 多少个 span (每个span经过一次放大器)
    span_length = 80  # 一个 span 的长度 (km)
    delta_z = 1  # 单步步长 (km)
    alpha = 0.2  # dB/km
    beta2 = -21.6676e-24  # s^2/km
    # beta2 = 0
    gamma = 1.3  # 1/(W*km)
    npol = 2

    signals, signals_power, f_omega, Beta2 =phot.fiber_G652(signals, num_symbols, bits_per_symbol, total_baud, up_sampling_factor, sampling_rate, npol, num_spans, span_length, delta_z, alpha, beta2, gamma)


    """" 相干接收机本振激光器 """
    LO_Linewidth = 100e3  # 激光器线宽
    frequency_offset = 2e9  # 频偏
    LO_CENTER_Frequency = Reference_frequency + frequency_offset  # LO中心频率
    LO_power_dB = 20  # power OF LO
    Optical_LO = phot.laser_cw(num_symbols * up_sampling_factor, LO_CENTER_Frequency, Reference_frequency,
                               sampling_rate, LO_power_dB, LO_Linewidth)
    # CD
    signals = phot.cd_compensation(input=signals, dispersion=17, Lambda=1550, FiberLength=num_spans * span_length,
                                   samplerate=sampling_rate, s=f_omega, f=Beta2)
    """" 相干接收机 """
    Responsivity = 1
    signals = phot.coherent_receiver(
        signals, Optical_LO, Responsivity)  # 由接收机组件组成的接收机

    # """ 接收端滤波器 """
    # filter_type = 'Gaussian'
    # BW_electrical = total_baud * 0.75
    # order = 4
    # signals = phot.Electrical_filter_POL(
    #     signals, sampling_rate, filter_type, BW_electrical, order)

    """ 模拟接收机造成的I/Q失衡，主要考虑幅度失衡和相位失衡，这里将两者都加在虚部上 """
    signals = phot.add_iq_imbalance(signals)  # 接收机的损伤

    """ 加入ADC的量化噪声 """
    adc_sample_rate = 120e9  # ADC采样率
    adc_resolution_bits = 8  # ADC的bit位数
    signals = phot.adc_noise(signals, sampling_rate,
                             adc_sample_rate, adc_resolution_bits)

    """ GSOP """
    signals = phot.iq_compensation(signals)

    """ 粗频偏估计和补偿 """
    signals = phot.freq_offset_compensation(signals, sampling_rate)

    signals = phot.pulse_shaper(
        signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    # 同步 #
    signals, prev_symbols = phot.synchronization(
        signals, prev_symbols, up_sampling_factor)

    """ 自适应均衡/解偏振：采用CMA-RDE算法 """
    num_tap = 25  # 均衡器抽头数目，此处均衡器内部是采用FIR滤波器，具体可查阅百度或者论文，
    ref_power_cma = 2  # 设置CMA算法的模
    cma_convergence = 30000  # CMA预均衡收敛的信号长度
    step_size_cma = 1e-9  # CMA的更新步长，梯度下降法的步长
    step_size_rde = 1e-9  # RDE的更新步长，梯度下降法的步长，%% CMA和RDE主要就是损失函数不同
    signals = phot.adaptive_equalizer(
        signals,
        num_tap,
        cma_convergence,
        ref_power_cma,
        step_size_cma,
        step_size_rde,
        up_sampling_factor,
        bits_per_symbol,
        total_baud,
    )
    phot.constellation_diagram(signals)

    # 精确的频偏估计和补偿：采用FFT-FOE算法，
    signals = phot.freq_offset_compensation(signals, total_baud)

    """ 相位恢复：采用盲相位搜索算法进行相位估计和补偿 """
    N_Test_Angle = 64  # BPS算法的测试角数目
    Block_Size = 100  # BPS算法的块长设置
    signals = phot.bps_restore(
        signals, N_Test_Angle, Block_Size, bits_per_symbol)

    phot.constellation_diagram(signals, isdata=False)
